import time
import os
import numpy as np
import matplotlib.pyplot as plt
import psutil
import json
from datetime import datetime

class PerformanceAnalyzer:
    """
    性能分析和可视化类
    用于分析和展示加密算法的性能指标
    """
    
    def __init__(self, results_dir='./results'):
        self.results_dir = results_dir
        os.makedirs(results_dir, exist_ok=True)
        self.results = []
    
    def benchmark_encryption(self, encrypt_func, data_sizes, key, iterations=3):
        """
        基准测试加密性能
        :param encrypt_func: 加密函数，接受(data, key)参数
        :param data_sizes: 测试数据大小列表（字节）
        :param key: 加密密钥
        :param iterations: 每个大小重复测试的次数
        :return: 性能结果字典
        """
        results = []
        
        for size in data_sizes:
            # 生成随机数据
            data = os.urandom(size)
            size_mb = size / (1024 * 1024)  # 转换为MB
            
            # 测量多次并计算平均值
            times = []
            memory_usages = []
            
            for _ in range(iterations):
                # 测量内存使用前的基准
                process = psutil.Process(os.getpid())
                baseline_memory = process.memory_info().rss / (1024 * 1024)  # MB
                
                # 测量加密时间
                start_time = time.time()
                encrypt_func(data, key)
                encryption_time = time.time() - start_time
                
                # 测量加密后的内存使用
                post_memory = process.memory_info().rss / (1024 * 1024)  # MB
                memory_delta = max(0, post_memory - baseline_memory)  # 防止负值
                
                times.append(encryption_time)
                memory_usages.append(memory_delta)
            
            # 计算统计值
            avg_time = np.mean(times)
            std_time = np.std(times)
            avg_memory = np.mean(memory_usages)
            
            # 计算吞吐量 (MB/s)
            throughput = size_mb / avg_time if avg_time > 0 else 0
            
            result = {
                'data_size_bytes': size,
                'data_size_mb': size_mb,
                'avg_encryption_time_s': avg_time,
                'std_encryption_time_s': std_time,
                'throughput_mbps': throughput,
                'avg_memory_usage_mb': avg_memory
            }
            
            results.append(result)
        
        # 保存结果
        self.results.extend(results)
        return results
    
    def benchmark_encryption_vs_decryption(self, encrypt_func, decrypt_func, data_sizes, key, iterations=3):
        """
        比较加密和解密性能
        :param encrypt_func: 加密函数，接受(data, key)参数
        :param decrypt_func: 解密函数，接受(encrypted_data, key)参数
        :param data_sizes: 测试数据大小列表（字节）
        :param key: 加密/解密密钥
        :param iterations: 每个大小重复测试的次数
        :return: 性能结果字典
        """
        results = []
        
        for size in data_sizes:
            # 生成随机数据
            data = os.urandom(size)
            size_mb = size / (1024 * 1024)  # 转换为MB
            
            # 预先加密数据以便解密测试
            encrypted_data = encrypt_func(data, key)
            
            # 测量多次并计算平均值
            encryption_times = []
            decryption_times = []
            
            for _ in range(iterations):
                # 测量加密时间
                start_time = time.time()
                encrypt_func(data, key)
                encryption_times.append(time.time() - start_time)
                
                # 测量解密时间
                start_time = time.time()
                decrypt_func(encrypted_data, key)
                decryption_times.append(time.time() - start_time)
            
            # 计算统计值
            avg_encryption_time = np.mean(encryption_times)
            avg_decryption_time = np.mean(decryption_times)
            
            # 计算吞吐量 (MB/s)
            encryption_throughput = size_mb / avg_encryption_time if avg_encryption_time > 0 else 0
            decryption_throughput = size_mb / avg_decryption_time if avg_decryption_time > 0 else 0
            
            result = {
                'data_size_bytes': size,
                'data_size_mb': size_mb,
                'avg_encryption_time_s': avg_encryption_time,
                'avg_decryption_time_s': avg_decryption_time,
                'encryption_throughput_mbps': encryption_throughput,
                'decryption_throughput_mbps': decryption_throughput,
                'time_ratio': avg_decryption_time / avg_encryption_time if avg_encryption_time > 0 else 0
            }
            
            results.append(result)
        
        return results
    
    def benchmark_key_generation(self, key_gen_func, key_sizes=None, iterations=10):
        """
        测试密钥生成性能
        :param key_gen_func: 密钥生成函数，接受key_size参数
        :param key_sizes: 密钥大小列表（位），如果为None则使用[128, 192, 256]
        :param iterations: 每个大小重复测试的次数
        :return: 性能结果字典
        """
        if key_sizes is None:
            key_sizes = [128, 192, 256]  # 默认密钥大小（位）
        
        results = []
        
        for key_size in key_sizes:
            key_size_bytes = key_size // 8  # 转换为字节
            
            # 测量多次并计算平均值
            times = []
            for _ in range(iterations):
                start_time = time.time()
                key_gen_func(key_size_bytes)
                key_gen_time = time.time() - start_time
                times.append(key_gen_time * 1000)  # 转换为毫秒
            
            # 计算统计值
            avg_time = np.mean(times)
            std_time = np.std(times)
            
            result = {
                'key_size_bits': key_size,
                'avg_generation_time_ms': avg_time,
                'std_generation_time_ms': std_time,
                'meets_requirement': avg_time <= 100  # 要求密钥生成时间≤100ms
            }
            
            results.append(result)
        
        return results
    
    def benchmark_memory_usage(self, encrypt_func, data_sizes, key):
        """
        测试内存使用情况
        :param encrypt_func: 加密函数，接受(data, key)参数
        :param data_sizes: 测试数据大小列表（字节）
        :param key: 加密密钥
        :return: 内存使用结果字典
        """
        results = []
        
        for size in data_sizes:
            # 生成随机数据
            data = os.urandom(size)
            size_mb = size / (1024 * 1024)  # 转换为MB
            
            # 测量内存使用前的基准
            process = psutil.Process(os.getpid())
            baseline_memory = process.memory_info().rss / (1024 * 1024)  # MB
            
            # 加密数据
            encrypt_func(data, key)
            
            # 测量加密后的内存使用
            post_memory = process.memory_info().rss / (1024 * 1024)  # MB
            memory_delta = max(0, post_memory - baseline_memory)  # 防止负值
            
            result = {
                'data_size_bytes': size,
                'data_size_mb': size_mb,
                'baseline_memory_mb': baseline_memory,
                'post_memory_mb': post_memory,
                'memory_delta_mb': memory_delta,
                'meets_requirement': memory_delta <= 512  # 要求内存占用≤512KB
            }
            
            results.append(result)
        
        return results
    
    def plot_throughput_vs_size(self, results=None):
        """
        绘制数据大小vs吞吐量图表
        :param results: 性能结果字典列表，如果为None则使用存储的结果
        :return: matplotlib figure对象
        """
        if results is None:
            results = self.results
        if not results:
            raise ValueError("No performance results available")
        
        # 提取数据
        sizes = [r['data_size_mb'] for r in results]
        throughputs = [r['throughput_mbps'] for r in results]
        
        # 创建图表
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.plot(sizes, throughputs, 'o-', linewidth=2)
        
        ax.set_title('加密吞吐量 vs 数据大小')
        ax.set_xlabel('数据大小 (MB)')
        ax.set_ylabel('吞吐量 (MB/s)')
        ax.grid(True)
        
        # 添加水平线表示性能要求（5MB/s）
        ax.axhline(y=5, color='r', linestyle='--', label='性能要求 (5 MB/s)')
        ax.legend()
        
        return fig
    
    def plot_encryption_vs_decryption(self, results):
        """
        绘制加密vs解密性能对比图
        :param results: 加密解密性能比较结果
        :return: matplotlib figure对象
        """
        if not results:
            raise ValueError("No comparison results available")
        
        # 提取数据
        sizes = [r['data_size_mb'] for r in results]
        enc_throughputs = [r['encryption_throughput_mbps'] for r in results]
        dec_throughputs = [r['decryption_throughput_mbps'] for r in results]
        
        # 创建图表
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.plot(sizes, enc_throughputs, 'bo-', linewidth=2, label='加密')
        ax.plot(sizes, dec_throughputs, 'go-', linewidth=2, label='解密')
        
        ax.set_title('加密 vs 解密性能')
        ax.set_xlabel('数据大小 (MB)')
        ax.set_ylabel('吞吐量 (MB/s)')
        ax.grid(True)
        ax.legend()
        
        return fig
    
    def plot_key_generation_time(self, results):
        """
        绘制密钥生成时间图表
        :param results: 密钥生成性能结果
        :return: matplotlib figure对象
        """
        if not results:
            raise ValueError("No key generation results available")
        
        # 提取数据
        key_sizes = [r['key_size_bits'] for r in results]
        times = [r['avg_generation_time_ms'] for r in results]
        std_devs = [r['std_generation_time_ms'] for r in results]
        
        # 创建图表
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.errorbar(key_sizes, times, yerr=std_devs, fmt='o-', linewidth=2, capsize=5)
        
        ax.set_title('密钥生成时间 vs 密钥大小')
        ax.set_xlabel('密钥大小 (bits)')
        ax.set_ylabel('生成时间 (ms)')
        ax.grid(True)
        
        # 添加水平线表示性能要求（100ms）
        ax.axhline(y=100, color='r', linestyle='--', label='性能要求 (100 ms)')
        ax.legend()
        
        return fig
    
    def plot_memory_usage(self, results):
        """
        绘制内存使用情况图表
        :param results: 内存使用结果
        :return: matplotlib figure对象
        """
        if not results:
            raise ValueError("No memory usage results available")
        
        # 提取数据
        sizes = [r['data_size_mb'] for r in results]
        memory_usage = [r['memory_delta_mb'] for r in results]
        
        # 创建图表
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.plot(sizes, memory_usage, 'o-', linewidth=2)
        
        ax.set_title('内存使用 vs 数据大小')
        ax.set_xlabel('数据大小 (MB)')
        ax.set_ylabel('内存增加 (MB)')
        ax.grid(True)
        
        # 添加水平线表示性能要求（0.5MB）
        ax.axhline(y=0.5, color='r', linestyle='--', label='性能要求 (512 KB)')
        ax.legend()
        
        return fig
    
    def save_results(self, filename=None):
        """
        保存性能结果到JSON文件
        :param filename: 文件名，如果为None则使用时间戳
        :return: 保存的文件路径
        """
        if not self.results:
            raise ValueError("No results to save")
        
        if filename is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"performance_results_{timestamp}.json"
        
        filepath = os.path.join(self.results_dir, filename)
        
        with open(filepath, 'w') as f:
            json.dump(self.results, f, indent=2)
        
        return filepath
    
    def save_plot(self, fig, filename=None):
        """
        保存图表到文件
        :param fig: matplotlib figure对象
        :param filename: 文件名，如果为None则使用时间戳
        :return: 保存的文件路径
        """
        if filename is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"performance_plot_{timestamp}.png"
        
        filepath = os.path.join(self.results_dir, filename)
        fig.savefig(filepath, dpi=300, bbox_inches='tight')
        
        return filepath
    
    def generate_performance_report(self, encrypt_func, decrypt_func, key_gen_func, key):
        """
        生成全面的性能报告
        :param encrypt_func: 加密函数
        :param decrypt_func: 解密函数
        :param key_gen_func: 密钥生成函数
        :param key: 测试用密钥
        :return: 报告文件路径
        """
        # 创建报告目录
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        report_dir = os.path.join(self.results_dir, f"report_{timestamp}")
        os.makedirs(report_dir, exist_ok=True)
        
        # 执行各项基准测试
        data_sizes = [1024, 1024*10, 1024*100, 1024*1024, 5*1024*1024, 10*1024*1024]  # 1KB到10MB
        
        print("测试加密性能...")
        enc_results = self.benchmark_encryption(encrypt_func, data_sizes, key)
        
        print("比较加密与解密性能...")
        enc_dec_results = self.benchmark_encryption_vs_decryption(encrypt_func, decrypt_func, data_sizes, key)
        
        print("测试密钥生成性能...")
        key_gen_results = self.benchmark_key_generation(key_gen_func)
        
        print("测试内存使用情况...")
        memory_results = self.benchmark_memory_usage(encrypt_func, data_sizes, key)
        
        # 生成并保存图表
        print("生成性能图表...")
        
        # 吞吐量图表
        fig1 = self.plot_throughput_vs_size(enc_results)
        throughput_plot_path = os.path.join(report_dir, "throughput_vs_size.png")
        fig1.savefig(throughput_plot_path, dpi=300, bbox_inches='tight')
        plt.close(fig1)
        
        # 加密vs解密图表
        fig2 = self.plot_encryption_vs_decryption(enc_dec_results)
        enc_dec_plot_path = os.path.join(report_dir, "encryption_vs_decryption.png")
        fig2.savefig(enc_dec_plot_path, dpi=300, bbox_inches='tight')
        plt.close(fig2)
        
        # 密钥生成时间图表
        fig3 = self.plot_key_generation_time(key_gen_results)
        key_gen_plot_path = os.path.join(report_dir, "key_generation_time.png")
        fig3.savefig(key_gen_plot_path, dpi=300, bbox_inches='tight')
        plt.close(fig3)
        
        # 内存使用图表
        fig4 = self.plot_memory_usage(memory_results)
        memory_plot_path = os.path.join(report_dir, "memory_usage.png")
        fig4.savefig(memory_plot_path, dpi=300, bbox_inches='tight')
        plt.close(fig4)
        
        # 保存原始数据
        print("保存性能数据...")
        with open(os.path.join(report_dir, "encryption_results.json"), 'w') as f:
            json.dump(enc_results, f, indent=2)
        
        with open(os.path.join(report_dir, "enc_dec_comparison.json"), 'w') as f:
            json.dump(enc_dec_results, f, indent=2)
        
        with open(os.path.join(report_dir, "key_generation_results.json"), 'w') as f:
            json.dump(key_gen_results, f, indent=2)
        
        with open(os.path.join(report_dir, "memory_results.json"), 'w') as f:
            json.dump(memory_results, f, indent=2)
        
        # 生成HTML报告
        print("生成HTML报告...")
        html_report_path = os.path.join(report_dir, "performance_report.html")
        self._generate_html_report(html_report_path, enc_results, enc_dec_results, key_gen_results, memory_results)
        
        print(f"性能报告已生成于: {report_dir}")
        return report_dir
    
    def _generate_html_report(self, filepath, enc_results, enc_dec_results, key_gen_results, memory_results):
        """
        生成HTML性能报告
        """
        # 检查是否满足性能要求
        meets_throughput = any(r['throughput_mbps'] >= 5 for r in enc_results)
        meets_key_gen = any(r['meets_requirement'] for r in key_gen_results)
        meets_memory = any(r['meets_requirement'] for r in memory_results)
        
        html_content = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <meta charset="UTF-8">
            <title>加密算法性能报告</title>
            <style>
                body {{ font-family: Arial, sans-serif; margin: 20px; }}
                h1, h2, h3 {{ color: #333; }}
                .report-section {{ margin-bottom: 30px; }}
                table {{ border-collapse: collapse; width: 100%; margin-bottom: 20px; }}
                th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
                th {{ background-color: #f2f2f2; }}
                tr:nth-child(even) {{ background-color: #f9f9f9; }}
                .pass {{ color: green; font-weight: bold; }}
                .fail {{ color: red; font-weight: bold; }}
                .chart {{ margin: 20px 0; text-align: center; }}
                .chart img {{ max-width: 800px; border: 1px solid #ddd; }}
                .summary {{ background-color: #f0f0f0; padding: 15px; border-radius: 5px; }}
            </style>
        </head>
        <body>
            <h1>加密算法性能报告</h1>
            <p>生成时间: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}</p>
            
            <div class="summary report-section">
                <h2>性能要求满足情况</h2>
                <ul>
                    <li>文本加密处理速度≥5MB/s: <span class="{'pass' if meets_throughput else 'fail'}">{meets_throughput}</span></li>
                    <li>密钥生成时间≤100ms: <span class="{'pass' if meets_key_gen else 'fail'}">{meets_key_gen}</span></li>
                    <li>内存占用≤512KB: <span class="{'pass' if meets_memory else 'fail'}">{meets_memory}</span></li>
                </ul>
            </div>
            
            <div class="report-section">
                <h2>加密性能</h2>
                <table>
                    <tr>
                        <th>数据大小 (MB)</th>
                        <th>加密时间 (秒)</th>
                        <th>吞吐量 (MB/s)</th>
                        <th>内存使用 (MB)</th>
                    </tr>
        """
        
        for r in enc_results:
            html_content += f"""
                    <tr>
                        <td>{r['data_size_mb']:.4f}</td>
                        <td>{r['avg_encryption_time_s']:.4f}</td>
                        <td>{r['throughput_mbps']:.4f}</td>
                        <td>{r['avg_memory_usage_mb']:.4f}</td>
                    </tr>
            """
        
        html_content += """
                </table>
                <div class="chart">
                    <img src="throughput_vs_size.png" alt="Throughput vs Size Chart">
                </div>
            </div>
            
            <div class="report-section">
                <h2>加密与解密性能比较</h2>
                <table>
                    <tr>
                        <th>数据大小 (MB)</th>
                        <th>加密时间 (秒)</th>
                        <th>解密时间 (秒)</th>
                        <th>加密吞吐量 (MB/s)</th>
                        <th>解密吞吐量 (MB/s)</th>
                        <th>时间比率 (解密/加密)</th>
                    </tr>
        """
        
        for r in enc_dec_results:
            html_content += f"""
                    <tr>
                        <td>{r['data_size_mb']:.4f}</td>
                        <td>{r['avg_encryption_time_s']:.4f}</td>
                        <td>{r['avg_decryption_time_s']:.4f}</td>
                        <td>{r['encryption_throughput_mbps']:.4f}</td>
                        <td>{r['decryption_throughput_mbps']:.4f}</td>
                        <td>{r['time_ratio']:.4f}</td>
                    </tr>
            """
        
        html_content += """
                </table>
                <div class="chart">
                    <img src="encryption_vs_decryption.png" alt="Encryption vs Decryption Chart">
                </div>
            </div>
            
            <div class="report-section">
                <h2>密钥生成性能</h2>
                <table>
                    <tr>
                        <th>密钥大小 (bits)</th>
                        <th>平均生成时间 (ms)</th>
                        <th>标准差 (ms)</th>
                        <th>满足要求 (≤100ms)</th>
                    </tr>
        """
        
        for r in key_gen_results:
            html_content += f"""
                    <tr>
                        <td>{r['key_size_bits']}</td>
                        <td>{r['avg_generation_time_ms']:.4f}</td>
                        <td>{r['std_generation_time_ms']:.4f}</td>
                        <td class="{'pass' if r['meets_requirement'] else 'fail'}">{r['meets_requirement']}</td>
                    </tr>
            """
        
        html_content += """
                </table>
                <div class="chart">
                    <img src="key_generation_time.png" alt="Key Generation Time Chart">
                </div>
            </div>
            
            <div class="report-section">
                <h2>内存使用情况</h2>
                <table>
                    <tr>
                        <th>数据大小 (MB)</th>
                        <th>基准内存 (MB)</th>
                        <th>加密后内存 (MB)</th>
                        <th>内存增加 (MB)</th>
                        <th>满足要求 (≤0.5MB)</th>
                    </tr>
        """
        
        for r in memory_results:
            html_content += f"""
                    <tr>
                        <td>{r['data_size_mb']:.4f}</td>
                        <td>{r['baseline_memory_mb']:.4f}</td>
                        <td>{r['post_memory_mb']:.4f}</td>
                        <td>{r['memory_delta_mb']:.4f}</td>
                        <td class="{'pass' if r['meets_requirement'] else 'fail'}">{r['meets_requirement']}</td>
                    </tr>
            """
        
        html_content += """
                </table>
                <div class="chart">
                    <img src="memory_usage.png" alt="Memory Usage Chart">
                </div>
            </div>
            
            <div class="report-section">
                <h2>结论</h2>
                <p>
                    基于以上性能测试结果，该加密算法实现
        """
        
        if meets_throughput and meets_key_gen and meets_memory:
            html_content += " <span class='pass'>满足</span> 所有性能要求。"
        else:
            html_content += " <span class='fail'>不满足</span> 全部性能要求。以下是详细情况："
            html_content += "<ul>"
            if not meets_throughput:
                html_content += "<li>加密速度未达到5MB/s的目标</li>"
            if not meets_key_gen:
                html_content += "<li>密钥生成时间超过100ms</li>"
            if not meets_memory:
                html_content += "<li>内存占用超过512KB</li>"
            html_content += "</ul>"
        
        html_content += """
                </p>
            </div>
        </body>
        </html>
        """
        
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(html_content)
        
        return filepath
    
    def analyze_text_operations(self, test_text, key_id=None, iterations=3, key_manager=None, crypto_engine=None):
        """
        分析文本操作（加密、解密、密钥检索）的性能
        
        :param test_text: 要测试的文本
        :param key_id: 密钥ID
        :param iterations: 每个操作重复测试的次数
        :param key_manager: 密钥管理器实例
        :param crypto_engine: 加密引擎实例
        :return: 性能测试结果
        """
        if not key_manager or not crypto_engine:
            raise ValueError("必须提供key_manager和crypto_engine")
        
        # 将文本转换为二进制
        text_data = test_text.encode('utf-8') if isinstance(test_text, str) else test_text
        
        # 获取当前系统资源使用情况
        process = psutil.Process(os.getpid())
        cpu_percent = psutil.cpu_percent(interval=0.1)
        memory_used = process.memory_info().rss / (1024 * 1024)  # MB
        
        # 1. 测量密钥检索性能
        key_times = []
        for _ in range(iterations):
            start_time = time.time()
            key = key_manager.get_key(key_id=key_id)
            key_time = (time.time() - start_time) * 1000  # 转为毫秒
            key_times.append(key_time)
        
        key_retrieval_results = {
            'avg_time': np.mean(key_times),
            'min_time': np.min(key_times),
            'max_time': np.max(key_times),
            'std_dev': np.std(key_times)
        }
        
        # 2. 测量加密性能
        encryption_times = []
        for _ in range(iterations):
            start_time = time.time()
            encrypted_data, iv = crypto_engine.encrypt(text_data, key)
            encryption_time = (time.time() - start_time) * 1000  # 转为毫秒
            encryption_times.append(encryption_time)
        
        encryption_results = {
            'avg_time': np.mean(encryption_times),
            'min_time': np.min(encryption_times),
            'max_time': np.max(encryption_times),
            'std_dev': np.std(encryption_times)
        }
        
        # 3. 测量解密性能
        decryption_times = []
        for _ in range(iterations):
            start_time = time.time()
            crypto_engine.decrypt(encrypted_data, key, iv)
            decryption_time = (time.time() - start_time) * 1000  # 转为毫秒
            decryption_times.append(decryption_time)
        
        decryption_results = {
            'avg_time': np.mean(decryption_times),
            'min_time': np.min(decryption_times),
            'max_time': np.max(decryption_times),
            'std_dev': np.std(decryption_times)
        }
        
        # 汇总结果
        return {
            'key_retrieval': key_retrieval_results,
            'encryption': encryption_results,
            'decryption': decryption_results,
            'data_size_bytes': len(text_data),
            'data_size_kb': len(text_data) / 1024,
            'iterations': iterations,
            'system': {
                'cpu_percent': cpu_percent,
                'memory_used': memory_used
            }
        } 